import face_recognition
import cv2

# 读取视频并获取帧的长度
input_movie = cv2.VideoCapture("../video/Video2024.wmv")
movie_length = int(input_movie.get(cv2.CAP_PROP_FRAME_COUNT))

# 指定输出视频编码方式、格式
fourcc = cv2.VideoWriter_fourcc(*'XVID')
output_movie = cv2.VideoWriter('test.avi', fourcc, 10, (1352, 756))

# 载入特征图片，并编码，编码完成后加入已知的图像列表
load_image_one = face_recognition.load_image_file("../img/huge.jpg")
# print(face_recognition.face_encodings(load_image))
load_face_encoding_one = face_recognition.face_encodings(load_image_one)[0]


load_image_two = face_recognition.load_image_file("../img/liuyifei.png")
# print(face_recognition.face_encodings(load_image))
load_face_encoding_two = face_recognition.face_encodings(load_image_two)[0]

known_faces = [
    load_face_encoding_one,
    load_face_encoding_two
]

face_locations = []
face_encodings = []
frame_number = 0

while True:
    ret, frame = input_movie.read()
    frame_number += 1
    
    if not ret or frame_number == 30:
        break

    # 在OpenCV中，图像默认是以BGR（蓝绿红）顺序排列的，而在许多其他库和应用中，
    # 如PIL（Python Imaging Library）和Matplotlib，图像是以RGB（红绿蓝）顺序排列的。
    # 因此，通过将图像的通道顺序反转（即::-1），我们可以将BGR格式的图像转换为RGB格式。

    # rgb_frame = frame[:, :, ::-1] 会导致程序崩溃
    rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    # 获取该帧图片的人脸的位置信息，返回一个人脸盒子组成的列表
    face_locations = face_recognition.face_locations(rgb_frame)
    # 将已知位置的人脸编码，并返回编码列表
    face_encodings = face_recognition.face_encodings(rgb_frame, face_locations)

    face_names = []

    for face_encoding in face_encodings:
        match = face_recognition.compare_faces(known_faces, face_encoding, 0.5)

        name = None
        if match[0]:
            name = "huge"
        elif match[1]:
            name = "*****"
        face_names.append(name)

    for (top, right, bottom, left), name in zip(face_locations, face_names):
        if not name:
            continue

        cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255), 2)

        cv2.rectangle(frame, (left, bottom - 25), (right, bottom), (0, 0, 255), cv2.FILLED)
        font = cv2.FONT_HERSHEY_DUPLEX
        cv2.putText(frame, name, (left + 6, bottom - 6), font, 0.5, (255, 255, 255), 1)

        print("writing frame {} / {}".format(frame_number, movie_length))
        output_movie.write(frame)

input_movie.release()
cv2.destroyAllWindows()



